import zmq
from kvsimple import KVMsg


class Route:
    def __init__(self, socket, identity, subtree):
        self.socket = socket
        self.identity = identity
        self.subtree = subtree


def send_single(key, kvmsg, route):
    if kvmsg.key.startswith(route.subtree):
        route.socket.send(route.identity, zmq.SENDMORE)
        kvmsg.send(route.socket)


def main():
    ctx = zmq.Context()
    snapshot = ctx.socket(zmq.ROUTER)
    snapshot.bind("tcp://*:5556")
    publisher = ctx.socket(zmq.PUB)
    publisher.bind("tcp://*:5557")
    collector = ctx.socket(zmq.PULL)
    collector.bind("tcp://*:5558")

    sequence = 0
    kvmap = {}

    poller = zmq.Poller()
    poller.register(collector, zmq.POLLIN)
    poller.register(snapshot, zmq.POLLIN)
    while True:
        try:
            items = dict(poller.poll(1000))
        except KeyboardInterrupt:
            break
        if collector in items:
            kvmsg = KVMsg.recv(collector)
            sequence += 1
            kvmsg.sequence = sequence
            kvmsg.send(publisher)
            kvmsg.store(kvmap)
            print("I: publishing update %5d" % sequence)
        if snapshot in items:
            msg = snapshot.recv_multipart()
            identity, request, subtree = msg
            if request == b"ICANHAZ?":
                pass
            else:
                print("E: bad request, aborting")
                break

            route = Route(snapshot, identity, subtree)
            for k, v in kvmap.items():
                send_single(k, v, route)

            print("Sending state snapshot=%d" % sequence)
            snapshot.send(identity, zmq.SNDMORE)
            kvmsg = KVMsg(sequence)
            kvmsg.key = "KTHXBAI"
            kvmsg.body = ""
            kvmsg.send(snapshot)

    print("Interrupted %d messages handled" % sequence)


if __name__ == '__main__':
    main()
